import os
import json
import base64
from typing import List, Dict, Any, Optional, Union
from pathlib import Path
import dashscope
from dashscope import Generation
from dashscope import MultiModalConversation
import logging
from pdf_processor import PDFProcessor
from config import DASHSCOPE_API_KEY, TEXT_MODEL, MULTIMODAL_MODEL, MEMORY_DIR

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 设置API密钥
dashscope.api_key = DASHSCOPE_API_KEY

class AIEngine:
    """AI引擎，用于处理用户查询和管理对话记忆"""
    
    def __init__(self, pdf_processor: PDFProcessor = None):
        """初始化AI引擎
        
        Args:
            pdf_processor: PDF处理器实例，用于向量搜索
        """
        self.pdf_processor = pdf_processor or PDFProcessor()
        self.memory_dir = Path(MEMORY_DIR)
        self.memory_dir.mkdir(exist_ok=True, parents=True)
        self.memory_file = self.memory_dir / "conversation_memory.json"
        
        # 创建对话历史目录
        self.conversations_dir = self.memory_dir / "conversations"
        self.conversations_dir.mkdir(exist_ok=True)
        
        # 当前活动对话ID
        self.current_conversation_id = None
        
        # 对话上下文，保存最近20轮对话
        self.conversation_history = []
        self.max_history_turns = 20
        
        # 对话顺序计数器文件
        self.order_counter_file = self.memory_dir / "conversation_order_counter.json"
        if not self.order_counter_file.exists():
            with open(self.order_counter_file, 'w', encoding='utf-8') as f:
                json.dump({"counter": 0}, f)
        
        self.load_memory()
        
    def _get_next_conversation_order(self) -> int:
        """获取下一个对话顺序号"""
        try:
            with open(self.order_counter_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
                counter = data.get("counter", 0)
            
            # 增加计数器
            counter += 1
            
            # 保存更新后的计数器
            with open(self.order_counter_file, 'w', encoding='utf-8') as f:
                json.dump({"counter": counter}, f)
                
            return counter
        except Exception as e:
            logger.error(f"获取对话顺序号失败: {str(e)}")
            # 如果出错，返回当前时间戳作为备用
            import time
            return int(time.time())
    
    def load_memory(self):
        """加载对话记忆"""
        if self.memory_file.exists():
            try:
                with open(self.memory_file, 'r', encoding='utf-8') as f:
                    self.memory = json.load(f)
            except Exception as e:
                logger.error(f"加载记忆文件失败: {str(e)}")
                self.memory = {"facts": [], "corrections": []}
        else:
            self.memory = {"facts": [], "corrections": []}
    
    def save_memory(self):
        """保存对话记忆"""
        try:
            with open(self.memory_file, 'w', encoding='utf-8') as f:
                json.dump(self.memory, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.error(f"保存记忆文件失败: {str(e)}")
    
    def add_to_memory(self, fact: str, is_correction: bool = False):
        """添加事实到记忆
        
        Args:
            fact: 要记住的事实
            is_correction: 是否为纠正信息
        """
        if is_correction:
            self.memory["corrections"].append(fact)
        else:
            self.memory["facts"].append(fact)
        self.save_memory()
    
    def delete_from_memory(self, index: int, is_correction: bool = False) -> bool:
        """从记忆中删除事实
        
        Args:
            index: 要删除的事实索引
            is_correction: 是否为纠正信息
            
        Returns:
            删除是否成功
        """
        try:
            if is_correction:
                if 0 <= index < len(self.memory["corrections"]):
                    del self.memory["corrections"][index]
                    self.save_memory()
                    return True
            else:
                if 0 <= index < len(self.memory["facts"]):
                    del self.memory["facts"][index]
                    self.save_memory()
                    return True
            return False
        except Exception as e:
            logger.error(f"删除记忆时出错: {str(e)}")
            return False
    
    def get_memory_context(self) -> str:
        """获取记忆上下文字符串
        
        Returns:
            记忆上下文字符串
        """
        context = "以下是我已知的关于网络实验的事实:\n"
        for fact in self.memory["facts"]:
            context += f"- {fact}\n"
        
        if self.memory["corrections"]:
            context += "\n以下是之前的错误认知的纠正:\n"
            for correction in self.memory["corrections"]:
                context += f"- {correction}\n"
        
        return context
        
    def _format_answer(self, text: str) -> str:
        """格式化和清理AI回答
        
        Args:
            text: 原始回答文本
            
        Returns:
            格式化后的回答
        """
        import re
        
        # 移除多余的分隔符和标记
        text = text.replace('---', '')
        text = text.replace('###', '')
        
        # 清理多余的空格和换行
        text = text.replace('\n\n\n', '\n\n')
        
        # 清理Markdown标记
        text = text.replace('```', '')
        
        # 替换序号格式为更简洁的形式
        text = re.sub(r'(\d+)\.\s+\*\*(.+?)\*\*', r'\1. \2', text)
        
        # 清理多余的星号
        text = text.replace('**', '')
        text = text.replace('*', '')
        
        # 清理多余的括号和特殊字符
        text = text.replace('（', '(')
        text = text.replace('）', ')')
        
        # 完全重写格式化逻辑 - 使用更简单直接的方法
        
        # 步骤1：首先处理所有数字列表项，确保它们前后都有换行
        # 匹配所有数字列表项（如"1. ", "2. "等）
        list_pattern = r'(\d+\.\s+)'
        
        # 找到所有数字列表项的起始位置
        list_matches = list(re.finditer(list_pattern, text))
        
        # 如果没有找到列表项，直接进行段落处理
        if not list_matches:
            # 简单地按空行分段
            paragraphs = []
            for para in re.split(r'\n\s*\n', text):
                if para.strip():
                    paragraphs.append(para.strip())
            return '\n\n'.join(paragraphs)
        
        # 步骤2：如果有列表项，按列表项分割文本
        result_parts = []
        last_end = 0
        
        for match in list_matches:
            start = match.start()
            
            # 添加列表项前的文本（如果有）
            if start > last_end:
                pre_text = text[last_end:start].strip()
                if pre_text:
                    result_parts.append(pre_text)
            
            # 找到这个列表项的内容（直到下一个列表项或文本结束）
            next_start = len(text)
            for next_match in list_matches:
                if next_match.start() > start:
                    next_start = next_match.start()
                    break
            
            # 提取列表项内容
            list_item = text[start:next_start].strip()
            result_parts.append(list_item)
            
            last_end = next_start
        
        # 添加最后一个列表项后的文本（如果有）
        if last_end < len(text):
            post_text = text[last_end:].strip()
            if post_text:
                result_parts.append(post_text)
        
        # 步骤3：处理每个部分，确保列表项独立成行
        final_parts = []
        for part in result_parts:
            # 检查这部分是否以列表项开头
            if re.match(list_pattern, part):
                # 这是一个列表项，直接添加
                final_parts.append(part)
            else:
                # 这是普通文本，按段落处理
                for para in re.split(r'\n\s*\n', part):
                    if para.strip():
                        final_parts.append(para.strip())
        
        # 组合成最终文本，段落之间用单个换行符分隔（而不是空行）
        formatted_text = '\n'.join(final_parts)
        
        # 最后一步：确保所有列表项都单独成行
        # 在数字+点+空格后添加换行（如果后面跟着文本）
        formatted_text = re.sub(r'(\d+\.\s+)([^\n])', r'\1\n\2', formatted_text)
        
        # 强制处理特殊情况：如果列表项后面有文本，确保文本正确换行
        # 例如 "1. 项目标题 项目内容" 变成 "1. 项目标题\n项目内容"
        lines = formatted_text.split('\n')
        processed_lines = []
        
        for line in lines:
            line = line.strip()
            if not line:
                processed_lines.append('')
                continue
                
            # 检查是否是列表项开头
            list_match = re.match(r'^(\d+\.\s+)(.+)$', line)
            if list_match:
                list_num = list_match.group(1)  # 列表编号部分 (如 "1. ")
                content = list_match.group(2)   # 内容部分
                
                # 检查内容是否包含多个句子或段落
                if '. ' in content or ': ' in content:
                    # 将第一个句号或冒号后的内容放到新行
                    first_sentence_end = max(content.find('. '), content.find(': '))
                    if first_sentence_end > 0:
                        first_part = content[:first_sentence_end+1]
                        rest_part = content[first_sentence_end+2:]
                        if rest_part:
                            processed_lines.append(f"{list_num}{first_part}")
                            processed_lines.append(rest_part)
                            continue
                
                # 如果没有特殊处理，保持原样
                processed_lines.append(line)
            else:
                processed_lines.append(line)
        
        # 重新组合文本，使用更紧凑的格式
        formatted_text = '\n'.join(processed_lines)
        
        return formatted_text
    
    def create_new_conversation(self) -> str:
        """创建新的对话，返回对话ID"""
        import uuid
        import datetime
        
        # 生成唯一ID
        conversation_id = str(uuid.uuid4())
        
        # 创建对话数据结构
        conversation_data = {
            "id": conversation_id,
            "created_at": datetime.datetime.now().isoformat(),
            "updated_at": datetime.datetime.now().isoformat(),
            "title": f"对话 {datetime.datetime.now().strftime('%Y-%m-%d %H:%M')}",
            "messages": [],
            "original_order": self._get_next_conversation_order()
        }
        
        # 保存对话
        conversation_file = self.conversations_dir / f"{conversation_id}.json"
        with open(conversation_file, 'w', encoding='utf-8') as f:
            json.dump(conversation_data, f, ensure_ascii=False, indent=2)
        
        # 设置为当前对话
        self.current_conversation_id = conversation_id
        self.conversation_history = []
        
        return conversation_id
    
    def load_conversation(self, conversation_id: str) -> bool:
        """加载指定的对话"""
        conversation_file = self.conversations_dir / f"{conversation_id}.json"
        if not conversation_file.exists():
            return False
        
        try:
            with open(conversation_file, 'r', encoding='utf-8') as f:
                conversation_data = json.load(f)
            
            # 加载对话历史
            self.conversation_history = conversation_data.get("messages", [])
            self.current_conversation_id = conversation_id
            return True
        except Exception as e:
            logger.error(f"加载对话失败: {str(e)}")
            return False
    
    def get_conversations_list(self) -> List[Dict[str, Any]]:
        """获取所有对话的列表"""
        conversations = []
        
        for file_path in self.conversations_dir.glob("*.json"):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    conversation_data = json.load(f)
                
                # 获取对话消息
                messages = conversation_data.get("messages", [])
                
                # 提取对话内容以显示在列表中
                first_user_message = ""
                if messages:
                    for msg in messages:
                        if msg.get("role") == "user":
                            first_user_message = msg.get("content", "")
                            break
                
                # 如果找到了第一条用户消息，截取前20个字符作为标题
                title = conversation_data.get("title", "")
                if first_user_message:
                    # 截取前20个字符，如果超过20个字符则添加省略号
                    title_from_message = first_user_message[:20] + ("..." if len(first_user_message) > 20 else "")
                    # 只有在标题是默认格式时才替换
                    if "对话 " in title and len(title) <= 20:
                        title = title_from_message
                
                conversations.append({
                    "id": conversation_data.get("id"),
                    "title": title,
                    "created_at": conversation_data.get("created_at"),
                    "updated_at": conversation_data.get("updated_at"),
                    "message_count": len(messages),
                    "original_order": conversation_data.get("original_order", 0)  # 添加原始顺序字段
                })
            except Exception as e:
                logger.error(f"读取对话文件失败 {file_path}: {str(e)}")
        
        # 不在后端排序，让前端决定排序方式
        return conversations
    
    def update_conversation(self, user_message: str, ai_message: str) -> bool:
        """更新当前对话历史"""
        if not self.current_conversation_id:
            self.create_new_conversation()
        
        conversation_file = self.conversations_dir / f"{self.current_conversation_id}.json"
        if not conversation_file.exists():
            return False
        
        try:
            # 读取当前对话
            with open(conversation_file, 'r', encoding='utf-8') as f:
                conversation_data = json.load(f)
            
            # 添加新消息
            import datetime
            now = datetime.datetime.now().isoformat()
            
            # 添加用户消息
            user_message_data = {
                "role": "user",
                "content": user_message,
                "timestamp": now
            }
            
            # 添加AI消息
            ai_message_data = {
                "role": "assistant",
                "content": ai_message,
                "timestamp": now
            }
            
            # 更新对话数据
            conversation_data["messages"].append(user_message_data)
            conversation_data["messages"].append(ai_message_data)
            conversation_data["updated_at"] = now
            
            # 更新内存中的历史记录，保持最近20轮
            self.conversation_history.append(user_message_data)
            self.conversation_history.append(ai_message_data)
            if len(self.conversation_history) > self.max_history_turns * 2:  # 每轮有两条消息
                self.conversation_history = self.conversation_history[-self.max_history_turns*2:]
            
            # 保存更新后的对话
            with open(conversation_file, 'w', encoding='utf-8') as f:
                json.dump(conversation_data, f, ensure_ascii=False, indent=2)
            
            return True
        except Exception as e:
            logger.error(f"更新对话失败: {str(e)}")
            return False
    
    def text_query(self, query: str, search_docs: bool = True, conversation_id: str = None) -> Dict[str, Any]:
        """处理文本查询
        
        Args:
            query: 用户查询
            search_docs: 是否搜索文档
            conversation_id: 对话ID，如果提供则加载该对话
            
        Returns:
            包含回复的字典
        """
        try:
            # 处理对话ID
            if conversation_id:
                # 如果提供了对话ID，加载该对话
                if conversation_id == "new":
                    self.create_new_conversation()
                else:
                    self.load_conversation(conversation_id)
            elif not self.current_conversation_id:
                # 如果没有当前对话，创建一个新对话
                self.create_new_conversation()
            
            # 构建系统提示
            system_prompt = "你是一个专业的计算机网络实验助手，擅长解答网络实验相关问题。请提供清晰、简洁的回答，避免使用过多的标记符号，保持良好的段落分隔。当使用数字列表（如1. 2. 3.）时，请确保每个列表项都单独成行。"
            
            # 添加记忆上下文
            memory_context = self.get_memory_context()
            if memory_context:
                system_prompt += f"\n\n{memory_context}"
            
            # 如果需要，搜索相关文档
            doc_context = ""
            if search_docs:
                search_results = self.pdf_processor.search_similar(query, collection_name="both", limit=3)
                if search_results:
                    doc_context = "以下是与你的问题相关的实验资料内容:\n\n"
                    for i, result in enumerate(search_results):
                        doc_context += f"[文档 {i+1}] "
                        doc_context += f"来源: {result['type']}，"
                        doc_context += f"文件: {result['metadata']['filename']}，"
                        doc_context += f"页码: {result['metadata']['page']}\n"
                        doc_context += f"{result['text'][:500]}...\n\n"
            
            # 构建完整的提示
            messages = [
                {"role": "system", "content": system_prompt}
            ]
            
            if doc_context:
                messages.append({"role": "system", "content": doc_context})
            
            # 添加对话历史上下文
            if self.conversation_history:
                # 只取最近的20轮对话（每轮包含用户和AI的消息）
                history_messages = self.conversation_history[-self.max_history_turns*2:]
                for msg in history_messages:
                    messages.append({"role": msg["role"], "content": msg["content"]})
            
            # 添加当前用户查询
            messages.append({"role": "user", "content": query})
            
            # 调用通义千问API
            response = Generation.call(
                model=TEXT_MODEL,
                messages=messages,
                result_format='message'
            )
            
            if response.status_code == 200:
                raw_answer = response.output.choices[0].message.content
                # 格式化回答
                answer = self._format_answer(raw_answer)
                
                # 更新对话历史
                self.update_conversation(query, answer)
                
                return {
                    "success": True,
                    "answer": answer,
                    "searched_docs": bool(doc_context),
                    "doc_count": len(search_results) if search_docs and 'search_results' in locals() else 0,
                    "conversation_id": self.current_conversation_id
                }
            else:
                logger.error(f"API调用失败: {response.code} - {response.message}")
                return {
                    "success": False,
                    "error": f"API调用失败: {response.message}"
                }
        except Exception as e:
            logger.error(f"处理文本查询时出错: {str(e)}")
            return {
                "success": False,
                "error": f"处理查询时出错: {str(e)}"
            }
    
    def multimodal_query(self, query: str, image_paths: List[str] = None, conversation_id: str = None) -> Dict[str, Any]:
        """处理多模态查询（文本+图像）
        
        Args:
            query: 用户查询
            image_paths: 图像文件路径列表
            conversation_id: 对话ID，如果提供则加载该对话
            
        Returns:
            包含回复的字典
        """
        try:
            # 处理对话ID
            if conversation_id:
                # 如果提供了对话ID，加载该对话
                if conversation_id == "new":
                    self.create_new_conversation()
                else:
                    self.load_conversation(conversation_id)
            elif not self.current_conversation_id:
                # 如果没有当前对话，创建一个新对话
                self.create_new_conversation()
            
            # 构建系统提示
            system_prompt = "你是一个专业的计算机网络实验助手，擅长解答网络实验相关问题，包括分析网络配置、拓扑图和错误截图。请提供清晰、简洁的回答，避免使用过多的标记符号，保持良好的段落分隔。当使用数字列表（如1. 2. 3.）时，请确保每个列表项都单独成行。"
            
            # 添加记忆上下文
            memory_context = self.get_memory_context()
            if memory_context:
                system_prompt += f"\n\n{memory_context}"
            
            # 构建消息
            user_content = [{"text": query}]
            
            # 添加图像（如果有）
            if image_paths:
                for img_path in image_paths:
                    # 检查文件是否存在
                    if not Path(img_path).exists():
                        logger.warning(f"图像文件不存在: {img_path}")
                        continue
                    
                    try:
                        # 读取图像并转换为Base64
                        with open(img_path, "rb") as img_file:
                            image_data = img_file.read()
                            
                            # 获取图片格式
                            img_format = "jpeg"  # 默认格式
                            if img_path.lower().endswith(".png"):
                                img_format = "png"
                            elif img_path.lower().endswith(".gif"):
                                img_format = "gif"
                            elif img_path.lower().endswith(".webp"):
                                img_format = "webp"
                            
                            # 添加正确的data URI前缀
                            base64_image = base64.b64encode(image_data).decode("utf-8")
                            base64_image_with_prefix = f"data:image/{img_format};base64,{base64_image}"
                            
                            logger.info(f"图片格式: {img_format}, 数据长度: {len(base64_image)}")
                        
                        # 添加到用户消息中
                        user_content.append({
                            "image": base64_image_with_prefix
                        })
                        logger.info(f"成功添加图像: {img_path}, 大小: {len(image_data)} 字节")
                    except Exception as e:
                        logger.error(f"处理图像时出错: {img_path}, 错误: {str(e)}")
            
            # 构建基本消息
            messages = [
                {"role": "system", "content": system_prompt}
            ]
            
            # 添加对话历史上下文
            if self.conversation_history:
                # 只取最近的对话历史，但排除包含图片的消息（多模态API不支持在历史中包含图片）
                history_messages = []
                for msg in self.conversation_history[-self.max_history_turns*2:]:
                    if msg["role"] == "user" and isinstance(msg["content"], list):
                        # 这是一个可能包含图片的用户消息，只保留文本部分
                        text_parts = [item["text"] for item in msg["content"] if "text" in item]
                        if text_parts:
                            history_messages.append({"role": "user", "content": " ".join(text_parts)})
                    else:
                        # 这是一个普通消息，直接添加
                        history_messages.append(msg)
                
                # 添加历史消息
                messages.extend(history_messages)
            
            # 添加当前用户消息（包含图片）
            messages.append({"role": "user", "content": user_content})
            
            # 记录请求详情
            logger.info(f"多模态请求: 模型={MULTIMODAL_MODEL}, 图像数量={len(image_paths) if image_paths else 0}")
            
            # 调用通义千问多模态API
            response = MultiModalConversation.call(
                model=MULTIMODAL_MODEL,
                messages=messages
            )
            
            if response.status_code == 200:
                try:
                    # 添加详细的响应结构日志
                    logger.info(f"响应结构: {type(response.output)}")
                    logger.info(f"响应内容: {response.output}")
                    
                    # 检查响应结构并提取文本
                    choices = response.output.get('choices', [])
                    if not choices:
                        raise ValueError("响应中没有choices字段")
                    
                    message = choices[0].get('message', {})
                    if not message:
                        raise ValueError("响应中没有message字段")
                    
                    content = message.get('content', [])
                    if not content:
                        raise ValueError("响应中没有content字段")
                    
                    # 遍历内容找到文本部分
                    text_parts = []
                    for item in content:
                        if isinstance(item, dict) and 'text' in item:
                            text_parts.append(item['text'])
                        elif isinstance(item, str):
                            text_parts.append(item)
                    
                    if not text_parts:
                        raise ValueError("响应中没有文本内容")
                    
                    raw_answer = ' '.join(text_parts)
                    
                    # 清理和格式化文本
                    answer = self._format_answer(raw_answer)
                    
                    # 更新对话历史
                    # 为了正确存储，我们需要将用户消息简化为纯文本
                    user_text_only = query
                    if image_paths and len(image_paths) > 0:
                        user_text_only = f"{query} [附带{len(image_paths)}张图片]"
                    self.update_conversation(user_text_only, answer)
                    
                    logger.info(f"多模态API调用成功，返回内容长度: {len(answer)}")
                    return {
                        "success": True,
                        "answer": answer,
                        "image_count": len(image_paths) if image_paths else 0,
                        "conversation_id": self.current_conversation_id
                    }
                except Exception as e:
                    logger.error(f"解析多模态API响应时出错: {str(e)}")
                    # 尝试提取任何可能的文本内容
                    try:
                        if hasattr(response, 'output') and response.output:
                            output_str = str(response.output)
                            # 截取一部分作为回复
                            answer = f"API返回了响应，但解析失败。响应片段: {output_str[:200]}..."
                            return {
                                "success": True,
                                "answer": answer,
                                "image_count": len(image_paths) if image_paths else 0
                            }
                    except:
                        pass
                    
                    return {
                        "success": False,
                        "error": f"解析响应时出错: {str(e)}"
                    }
            else:
                logger.error(f"多模态API调用失败: {response.status_code} - {response.message if hasattr(response, 'message') else '未知错误'}")
                return {
                    "success": False,
                    "error": f"多模态API调用失败: {response.message if hasattr(response, 'message') else '未知错误'}"
                }
        except Exception as e:
            logger.error(f"处理多模态查询时出错: {str(e)}")
            return {
                "success": False,
                "error": f"处理查询时出错: {str(e)}"
            }
